#!/usr/bin/env python
# -*- coding: utf-8 -*-
# local proxy for video
# Based on GAppProxy by Du XiaoGang <dugang@188.com>

__author__ = 'd3d3LmVodXN0QGdtYWlsLmNvbQ=='.decode('base64')
__version__ = '0.3.8'

import BaseHTTPServer, SocketServer, socket, errno, re, time, util
import rangefetch
import copy
import thread
import traceback
from headermessage import headermessage
from Wcrypto import RandomMethod

try:
    import ssl
    ssl_enabled = True
except:
    ssl_enabled = False
class GetOut(Exception):
    pass

class LocalProxyHandler(BaseHTTPServer.BaseHTTPRequestHandler):
    #partSize = 1024000
    FR_Headers = ('', 'host', 'vary', 'via', 'x-forwarded-for',
            'proxy-authorization', 'proxy-connection', 'upgrade', 'keep-alive')

    def send_response(self, code, message=None):
        self.log_request(code)
        if message is None:
            if code in self.responses:
                message = self.responses[code][0]
            else:
                message = 'WallProxy Notify'
        if self.request_version != 'HTTP/0.9':
            self.wfile.write("%s %d %s\r\n" % (self.protocol_version, code, message))

    def end_error(self, code, message=None, data=None):
        if not data:
            self.send_error(code, message)
            
        else:
            self.send_response(code, message)
            self.wfile.write(data)
        self.connection.close()

    def do_CONNECT(self):
        if not ssl_enabled:
            return self.end_error(501, 'Local proxy error, HTTPS needs Python2.6 or later')

        # for ssl proxy
        host, _, port = self.path.rpartition(':')
        keyFile, crtFile = util.getCertificate(host)
        self.send_response(200)
        self.end_headers()
        try:
            ssl_sock = ssl.wrap_socket(self.connection, keyFile, crtFile, True)
        except ssl.SSLError, e:
            print 'SSLError: ' + str(e)
            return

        # rewrite request line, url to abs
        first_line = ''
        while True:
            data = ssl_sock.read()
            # EOF?
            if data == '':
                # bad request
                ssl_sock.close()
                self.connection.close()
                return
            # newline(\r\n)?
            first_line += data
            if '\n' in first_line:
                first_line, data = first_line.split('\n', 1)
                first_line = first_line.rstrip('\r')
                break
        # got path, rewrite
        method, path, ver = first_line.split()
        if path.startswith('/'):
            path = 'https://%s%s' % (host if port=='443' else self.path, path)
        # connect to local proxy server
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if not util.LISTEN_ADDR[0]:
            sock.connect(('127.0.0.1', util.LISTEN_ADDR[1]))
        else:
            sock.connect(util.LISTEN_ADDR)
        sock.send('%s %s %s\r\n%s' % (method, path, ver, data))

        # forward https request
        ssl_sock.settimeout(1)
        while True:
            try:
                data = ssl_sock.read(8192)
            except ssl.SSLError, e:
                if str(e).lower().find('timed out') == -1:
                    # error
                    sock.close()
                    ssl_sock.close()
                    self.connection.close()
                    return
                # timeout
                break
            if data != '':
                sock.send(data)
            else:
                # EOF
                break

        ssl_sock.setblocking(True)
        # simply forward response
        while True:
            data = sock.recv(8192)
            if data != '':
                ssl_sock.write(data)
            else:
                # EOF
                break
        # clean
        sock.close()
        ssl_sock.shutdown(socket.SHUT_WR)
        ssl_sock.close()
        self.connection.close()
        
    #add some module to do this
    _RangeFetch=rangefetch._RangeFetch
    
    def do_METHOD(self):
        try:
            self._do_METHOD()
        except socket.error, (err, _):
            # Connection closed before proxy return
            if err == errno.EPIPE or err == 10053: 
                print 'lost connet with',self.client_address[0]
        #except Exception, e:
        #    print 'UnExpected Error',repr(e)

    def _do_METHOD(self):
        partsize=util.RANGEFETCH_PARTSIZE
        if self.path.startswith('/'):
            host = self.headers['host']
            if host.endswith(':80'):
                host = host[:-3]
            self.path = 'http://%s%s' % (host , self.path)
        handler, proxy = util.findProxy(self.client_address[0], self.path, self.command, dict(self.headers))
        if handler is None:
            return self.end_error(*proxy)

        payload_len = int(self.headers.get('content-length', 0))
        if payload_len > 0:
            payload = self.rfile.read(payload_len)
        else:
            payload = ''

        for k in LocalProxyHandler.FR_Headers:
            try:
                del self.headers[k]
            except KeyError:
                pass
        
        CR_header = AddHeaderRange(self.path,self.headers,self.command,partsize)
        
        retval, data = handler(proxy, self.path, self.command, CR_header, payload)
        try:
            if retval == -1:
                return self.end_error(502, str(data))
            # Proxy Fetch
            if isinstance(data, dict):
                if data['code']==206 and self.command=='GET':
                    m = re.search(r'bytes\s+(\d+)-(\d+)/(\d+)', data['headers'].get('content-range',''))
                    if m and self._RangeFetch(handler, proxy, m, data):
                        return
                self.send_response(data['code'])
                
                for k,v in data['headers'].iteritems():
                    self.send_header(k.title(), v)
                self.end_headers()
                self.wfile.write(data['content'])
            # Direct Fetch
            else:
                resp = data
                self.send_response(resp.code)
                # Write resp.headers directly to avoid set-cookie bug
                self.wfile.write(resp.headers)
                self.end_headers()
                data = resp.read(8192)
                while data != '':
                    self.wfile.write(data)
                    data = resp.read(8192)
                resp.close()
        except socket.error, (err, _):
            # Connection closed before proxy return
            # just for Linux
            if err == errno.EPIPE or err == 10053: return
        self.connection.close()
    def log_message(self, format, *args):
        loginfo("%s - - [%s] %s" %
                         (self.address_string(),
                          self.log_date_time_string(),
                          format%args))
    do_GET = do_METHOD
    do_HEAD = do_METHOD
    do_PUT = do_METHOD
    do_POST = do_METHOD
    do_DELETE = do_METHOD
class ThreadingHTTPServer(SocketServer.ThreadingMixIn, BaseHTTPServer.HTTPServer):
    def handle_error(self, request, client_address):
        #override handle_error to avoid traceback when client lost connect
        formatted_lines = traceback.format_exc().splitlines()
        s='lose connect with ' + str(client_address) +'---'+ str(formatted_lines[-1])
        #logthis
        loginfo(s)

def AddHeaderRange(url,header,command,partsize):
        #try to add 'range' header
        #this action can reduce time to get first part data of video
        #some pages will go wrong
        #use function IsAddHeaderRange in proxy.conf to config this
    try:
        #Get a copy of header
        CR_header=headermessage(header)
        if (util.IsAddHeaderRange(url,command)==0): return CR_header 
        if ('range' in CR_header):
            pass
        else:   raise GetOut('')#cann't find 'range' header
        m=re.search(r'(\d+)?-(\d+)?',CR_header['range'])
        #range format may be wrong
        if not m: raise GetOut('')
        m = [u and int(u) for u in m.groups()]
        #range format like '-' no data both side
        if m[0] is None and m[1] is None: raise GetOut('')
        if m[0] is None:
        #range format like '-xx' no begin place
            if m[1] > partsize:
                m[1] = partsize
        elif m[1] is None or m[1]-m[0]+1 > partsize:
        #range format like ('xx-') or ('xx-xx' and override limit)
            m[1] = partsize - 1 + m[0]
        fetch_range = ('bytes=%s-%s' % (m[0], m[1])).replace('None', '')
    except GetOut , e:
        fetch_range = ('bytes=%s-%s' % (0 , partsize-1))
    except KeyError ,e:
        fetch_range = ('bytes=%s-%s' % (0 , partsize-1))
    CR_header['range']=fetch_range
    return CR_header
    
        
def funnothing(*x):pass
def printstr(s):print str(s)
if util.LOG=='logging':
  import logging
  LOG_FILENAME = 'proxy.log'
  logging.basicConfig(filename=LOG_FILENAME,level=logging.DEBUG)
  loginfo = logging.info
elif util.LOG=='stdout':
  loginfo = printstr
else:
#do nothing
  loginfo=funnothing

if __name__ == '__main__':
    print '---------------------------------------------------'
    print 'ListenAddress: %s:%d' % util.LISTEN_ADDR
    print 'HTTPS Enabled: %s' % ('YES' if ssl_enabled else 'NO')
    print 'OpenSSLModule: %s' % ('YES' if util.crypto else 'NO')
    print 'Proxies Count: GAE(%d) PHP(%d) DIRECT(%d)' % (len(util.GAE_PROXY), len(util.PHP_PROXY), len(util.DIRECT_PROXY))
    print 'Random Method: %s' % RandomMethod()
    print '---------------------------------------------------'
    httpd = ThreadingHTTPServer(util.LISTEN_ADDR, LocalProxyHandler)
    if util.RELOAD:
        from threading import Thread
        def repeat():
            while util.RELOAD:
                time.sleep(util.RELOAD)
                util.checkConf()
        Thread(target=repeat).start()
    #used to quit in linux shell,and release the port
    #you can alse quit in windows shell
    def endhttp(httpd):
        try:
            print "input 'q' try to quit in linux"
            while(True):
                a_input=raw_input()
                if (a_input=='q'): break
            httpd.shutdown()     
            util.RELOAD = 0
        except Exception,e:
            print 'UnExpected Error',repr(e)
    thread.start_new_thread(endhttp,tuple([httpd]))
    httpd.serve_forever()

